#-------------------------------------------------------------------------------
# Name:        OpenLibSys
# Purpose:     Windows device IO access from WinRing0.sys
#
# Author:      SUPER_紫电  QQ:911344755
#
# Created:     15/08/2022
# Copyright:   (c) SUPER_紫电 2022 All rights reserved.
# Licence:     <Modified BSD License>
#-------------------------------------------------------------------------------

import os
import sys
import platform
from ctypes import c_ubyte, c_ushort, c_ulong, c_void_p, sizeof, byref, pointer
import Driver as drv
import OlsIoctl as olsIoctl
from OlsDef import *
from win32 import *

BYTE = c_ubyte
WORD = c_ushort
DWORD = c_ulong

def is_64bit_os():
    if "PROCESSOR_ARCHITEW6432" in os.environ:
        if os.environ["PROCESSOR_ARCHITEW6432"] == "AMD64":
               return True
    if "64" in os.environ["PROCESSOR_ARCHITECTURE"]:
        return True
    return False

#-----------------------------------------------------------------------------
#
# Global
#
#-----------------------------------------------------------------------------

gPciMaxBus = 255
gPciMaxDevice = 31
gPciMaxFunction = 7

#gIsNT = False

#gDriverFileName = ""
#gDriverPath = ""

#-----------------------------------------------------------------------------
#
# Initialize/Deinitialize Functions
#
#-----------------------------------------------------------------------------
def get_dir():
    try:
        dirname = os.path.dirname(__file__)
    except:
        import inspect
        dirname = os.path.dirname(inspect.getsourcefile(lambda:0))
    return dirname

def Initialize():
    global gDriverFileName
    global gDriverPath
    if is_64bit_os():
        gDriverFileName = "WinRing0x64.sys"
    else:
        gDriverFileName = "WinRing0x32.sys"
    gDriverPath = os.path.join(
        #os.path.dirname(__file__),
        get_dir(),
        gDriverFileName)
    return os.path.exists(gDriverPath)

def InitOpenLibSys():
    global gIsNT
    if not Initialize():
        return False
    if platform.system() != "Windows":
        return False
    gIsNT = sys.getwindowsversion()[3] == 2
    if not gIsNT:
        return False
    if drv.OpenDriver(olsIoctl.OLS_DRIVER_ID):
        return True
    drv.ManageDriver(olsIoctl.OLS_DRIVER_ID, gDriverPath, \
        drv.DRIVER_REMOVE)
    if not drv.ManageDriver(olsIoctl.OLS_DRIVER_ID, gDriverPath, \
        drv.DRIVER_INSTALL):
        drv.ManageDriver(olsIoctl.OLS_DRIVER_ID, gDriverPath, \
            drv.DRIVER_REMOVE)
        return False
    return drv.OpenDriver(olsIoctl.OLS_DRIVER_ID)

def DeinitOpenLibSys():
    if gIsNT and GetRefCount() == 1:
        CloseHandle(drv.gHandle)
        drv.gHandle = INVALID_HANDLE_VALUE
        drv.ManageDriver(olsIoctl.OLS_DRIVER_ID, gDriverPath, \
            drv.DRIVER_REMOVE)
    if drv.gHandle != INVALID_HANDLE_VALUE:
        CloseHandle(drv.gHandle)
        drv.gHandle = INVALID_HANDLE_VALUE

def GetRefCount():
    if drv.gHandle == INVALID_HANDLE_VALUE:
        return 0
    outBuffer = DWORD(0)
    dwBytesReturned = DWORD(0)
    result = DeviceIoControl(
        drv.gHandle,
        olsIoctl.IOCTL_OLS_GET_REFCOUNT,
        None,
        0,
        byref(outBuffer),
        sizeof(outBuffer),
        byref(dwBytesReturned),
        None
    )
    return outBuffer.value if result else 0

def InvalidArgumentError(ValueError):
    pass

#-----------------------------------------------------------------------------
#
# I/O Read
#
#-----------------------------------------------------------------------------

def ReadIoPort(portAddr, dataSize):
    if dataSize not in [1, 2, 4]:
        raise InvalidArgumentError(
            'Argument "dataSize" only accept 1, 2, 4. Current value : %s!'
            % (str(dataSize))
        )
    if drv.gHandle == INVALID_HANDLE_VALUE:
        return 0
    if dataSize == 1:
        outBuffer = BYTE(0)
        ioControlCode = olsIoctl.IOCTL_OLS_READ_IO_PORT_BYTE
    elif dataSize == 2:
        outBuffer = WORD(0)
        ioControlCode = olsIoctl.IOCTL_OLS_READ_IO_PORT_WORD
    elif dataSize == 4:
        outBuffer = DWORD(0)
        ioControlCode = olsIoctl.IOCTL_OLS_READ_IO_PORT_DWORD
    inBuffer = DWORD(portAddr)
    dwBytesReturned = DWORD(0)
    status = DeviceIoControl(
        drv.gHandle,
        ioControlCode,
        byref(inBuffer),
        sizeof(inBuffer),
        byref(outBuffer),
        sizeof(outBuffer),
        byref(dwBytesReturned),
        None
    )
    return outBuffer.value, status

#-----------------------------------------------------------------------------
# I/O Read (No Error Check)
#-----------------------------------------------------------------------------

def ReadIoPortByte(port):
    #return ReadIoPort(port, 1)[0]
    if drv.gHandle == INVALID_HANDLE_VALUE:
        return 0
    outBuffer = BYTE(0)
    ioControlCode = olsIoctl.IOCTL_OLS_READ_IO_PORT_BYTE
    inBuffer = DWORD(port)
    dwBytesReturned = DWORD(0)
    status = DeviceIoControl(
        drv.gHandle,
        ioControlCode,
        byref(inBuffer),
        sizeof(inBuffer),
        byref(outBuffer),
        sizeof(outBuffer),
        byref(dwBytesReturned),
        None
    )
    return outBuffer.value

def ReadIoPortWord(port):
    #return ReadIoPort(port, 2)[0]
    if drv.gHandle == INVALID_HANDLE_VALUE:
        return 0
    outBuffer = WORD(0)
    ioControlCode = olsIoctl.IOCTL_OLS_READ_IO_PORT_WORD
    inBuffer = DWORD(port)
    dwBytesReturned = DWORD(0)
    status = DeviceIoControl(
        drv.gHandle,
        ioControlCode,
        byref(inBuffer),
        sizeof(inBuffer),
        byref(outBuffer),
        sizeof(outBuffer),
        byref(dwBytesReturned),
        None
    )
    return outBuffer.value

def ReadIoPortDword(port):
    #return ReadIoPort(port, 4)[0]
    if drv.gHandle == INVALID_HANDLE_VALUE:
        return 0
    outBuffer = DWORD(0)
    ioControlCode = olsIoctl.IOCTL_OLS_READ_IO_PORT_DWORD
    inBuffer = DWORD(port)
    dwBytesReturned = DWORD(0)
    status = DeviceIoControl(
        drv.gHandle,
        ioControlCode,
        byref(inBuffer),
        sizeof(inBuffer),
        byref(outBuffer),
        sizeof(outBuffer),
        byref(dwBytesReturned),
        None
    )
    return outBuffer.value

#-----------------------------------------------------------------------------
# I/O Read (Error Check)
#-----------------------------------------------------------------------------

def ReadIoPortByteEx(port):
    #return ReadIoPort(port, 1)
    if drv.gHandle == INVALID_HANDLE_VALUE:
        return 0
    outBuffer = BYTE(0)
    ioControlCode = olsIoctl.IOCTL_OLS_READ_IO_PORT_BYTE
    inBuffer = DWORD(port)
    dwBytesReturned = DWORD(0)
    status = DeviceIoControl(
        drv.gHandle,
        ioControlCode,
        byref(inBuffer),
        sizeof(inBuffer),
        byref(outBuffer),
        sizeof(outBuffer),
        byref(dwBytesReturned),
        None
    )
    return outBuffer.value, status

def ReadIoPortWordEx(port):
    #return ReadIoPort(port, 2)
    if drv.gHandle == INVALID_HANDLE_VALUE:
        return 0
    outBuffer = WORD(0)
    ioControlCode = olsIoctl.IOCTL_OLS_READ_IO_PORT_WORD
    inBuffer = DWORD(port)
    dwBytesReturned = DWORD(0)
    status = DeviceIoControl(
        drv.gHandle,
        ioControlCode,
        byref(inBuffer),
        sizeof(inBuffer),
        byref(outBuffer),
        sizeof(outBuffer),
        byref(dwBytesReturned),
        None
    )
    return outBuffer.value, status

def ReadIoPortDwordEx(port):
    #return ReadIoPort(port, 4)
    if drv.gHandle == INVALID_HANDLE_VALUE:
        return 0
    outBuffer = DWORD(0)
    ioControlCode = olsIoctl.IOCTL_OLS_READ_IO_PORT_DWORD
    inBuffer = DWORD(port)
    dwBytesReturned = DWORD(0)
    status = DeviceIoControl(
        drv.gHandle,
        ioControlCode,
        byref(inBuffer),
        sizeof(inBuffer),
        byref(outBuffer),
        sizeof(outBuffer),
        byref(dwBytesReturned),
        None
    )
    return outBuffer.value, status

#-----------------------------------------------------------------------------
#
# I/O Write
#
#-----------------------------------------------------------------------------

def WriteIoPort(portAddr, dataValue, dataSize):
    if dataSize not in [1, 2, 4]:
        raise InvalidArgumentError(
            'Argument "dataSize" only accept 1, 2, 4. Current value : %s!'
            % (str(dataSize))
        )
    if drv.gHandle == INVALID_HANDLE_VALUE:
        return
    inBuffer = olsIoctl.OLS_WRITE_IO_PORT_INPUT()
    inBuffer.PortAddr = portAddr
    if dataSize == 1:
        inBuffer.DataValue.CharData = dataValue
        ioControlCode = olsIoctl.IOCTL_OLS_WRITE_IO_PORT_BYTE
    elif dataSize == 2:
        inBuffer.DataValue.ShortData = dataValue
        ioControlCode = olsIoctl.IOCTL_OLS_WRITE_IO_PORT_WORD
    elif dataSize == 4:
        inBuffer.DataValue.LongData = dataValue
        ioControlCode = olsIoctl.IOCTL_OLS_WRITE_IO_PORT_DWORD
    dwBytesReturned = DWORD(0)
    status = DeviceIoControl(
        drv.gHandle,
        ioControlCode,
        pointer(inBuffer),
        sizeof(inBuffer),
        None,
        0,
        byref(dwBytesReturned),
        None
    )
    return status

def WriteIoPortByte(port, value):
    #return WriteIoPort(port, value, 1)
    if drv.gHandle == INVALID_HANDLE_VALUE:
        return False
    inBuffer = olsIoctl.OLS_WRITE_IO_PORT_INPUT()
    inBuffer.PortAddr = port
    inBuffer.DataValue.CharData = value
    ioControlCode = olsIoctl.IOCTL_OLS_WRITE_IO_PORT_BYTE
    dwBytesReturned = DWORD(0)
    status = DeviceIoControl(
        drv.gHandle,
        ioControlCode,
        pointer(inBuffer),
        sizeof(inBuffer),
        None,
        0,
        byref(dwBytesReturned),
        None
    )
    return status

def WriteIoPortWord(port, value):
    #return WriteIoPort(port, value, 2)
    if drv.gHandle == INVALID_HANDLE_VALUE:
        return False
    inBuffer = olsIoctl.OLS_WRITE_IO_PORT_INPUT()
    inBuffer.PortAddr = port
    inBuffer.DataValue.ShortData = value
    ioControlCode = olsIoctl.IOCTL_OLS_WRITE_IO_PORT_WORD
    dwBytesReturned = DWORD(0)
    status = DeviceIoControl(
        drv.gHandle,
        ioControlCode,
        pointer(inBuffer),
        sizeof(inBuffer),
        None,
        0,
        byref(dwBytesReturned),
        None
    )
    return status

def WriteIoPortDword(port, value):
    #return WriteIoPort(port, value, 4)
    if drv.gHandle == INVALID_HANDLE_VALUE:
        return False
    inBuffer = olsIoctl.OLS_WRITE_IO_PORT_INPUT()
    inBuffer.PortAddr = port
    inBuffer.DataValue.LongData = value
    ioControlCode = olsIoctl.IOCTL_OLS_WRITE_IO_PORT_DWORD
    dwBytesReturned = DWORD(0)
    status = DeviceIoControl(
        drv.gHandle,
        ioControlCode,
        pointer(inBuffer),
        sizeof(inBuffer),
        None,
        0,
        byref(dwBytesReturned),
        None
    )
    return status

#-----------------------------------------------------------------------------
#
# PCI Configuration Read
#
#-----------------------------------------------------------------------------

def ReadPciConfig(pciAddress, regAddress, lpOutBuffer, dataSize):
    if drv.gHandle == INVALID_HANDLE_VALUE:
        return False
    if lpOutBuffer == None:
        return False
    # alignment check
    if(dataSize == 2 and (regAddress & 1) != 0):
        return False
    if(dataSize == 4 and (regAddress & 3) != 0):
        return False
    inBuffer = olsIoctl.OLS_READ_PCI_CONFIG_INPUT()
    inBuffer.PciAddress = pciAddress
    inBuffer.PciOffset = regAddress
    dwBytesReturned = DWORD(0)
    status = DeviceIoControl(
        drv.gHandle,
        olsIoctl.IOCTL_OLS_READ_PCI_CONFIG,
        byref(inBuffer),
        sizeof(inBuffer),
        lpOutBuffer,
        dataSize,
        byref(dwBytesReturned),
        None
    )
    return status

#-----------------------------------------------------------------------------
# PCI Configuration Read (Return Value, Status)
#-----------------------------------------------------------------------------

def ReadPciConfigByteEx(pciAddress, regAddress):
    data = BYTE(0)
    if ReadPciConfig(pciAddress, regAddress, byref(data), 1):
        return data.value & 0xFF, True
    else:
        return 0xFF, False

def ReadPciConfigWordEx(pciAddress, regAddress):
    data = WORD(0)
    if ReadPciConfig(pciAddress, regAddress, byref(data), 2):
        return data.value, True
    else:
        return 0xFFFF, False

def ReadPciConfigDwordEx(pciAddress, regAddress):
    data = DWORD(0)
    if ReadPciConfig(pciAddress, regAddress, byref(data), 4):
        return data.value, True
    else:
        return 0xFFFFFFFF, False

#-----------------------------------------------------------------------------
# PCI Configuration Read (Return Value)
#-----------------------------------------------------------------------------

def ReadPciConfigByte(pciAddress, regAddress):
    return ReadPciConfigByteEx(pciAddress, regAddress)[0]

def ReadPciConfigWord(pciAddress, regAddress):
    return ReadPciConfigWordEx(pciAddress, regAddress)[0]

def ReadPciConfigDword(pciAddress, regAddress):
    return ReadPciConfigDwordEx(pciAddress, regAddress)[0]

#-----------------------------------------------------------------------------
#
# PCI Configuration Write
#
#-----------------------------------------------------------------------------

def WritePciConfig(pciAddress, regAddress, dataList, dataSize):
    if drv.gHandle == INVALID_HANDLE_VALUE:
        return False
    # alignment check
    if(dataSize == 2 and (regAddress & 1) != 0):
        return False
    if(dataSize == 4 and (regAddress & 3) != 0):
        return False
    inBuffer = [pciAddress, regAddress]
    inBuffer.extend(dataList)
    dwBytesReturned = DWORD(0)
    status = DeviceIoControl(
        drv.gHandle,
        olsIoctl.IOCTL_OLS_WRITE_PCI_CONFIG,
        byref(inBuffer),
        sizeof(inBuffer),
        None,
        0,
        byref(dwBytesReturned),
        None
    )
    return status

#-----------------------------------------------------------------------------
# PCI Configuration Write (Error Check)
#-----------------------------------------------------------------------------

def WritePciConfigByte(pciAddress, regAddress, dataValue):
    dataList = [BYTE(dataValue)]
    return WritePciConfig(pciAddress, regAddress, dataList, 1)

def WritePciConfigWord(pciAddress, regAddress, dataValue):
    dataList = [WORD(dataValue)]
    return WritePciConfig(pciAddress, regAddress, dataList, 2)

def WritePciConfigDword(pciAddress, regAddress, dataValue):
    dataList = [DWORD(dataValue)]
    return WritePciConfig(pciAddress, regAddress, dataList, 4)

#-----------------------------------------------------------------------------
#
# FindPciDevice
#
#-----------------------------------------------------------------------------

def SetPciMaxBusIndex(_max):
    global gPciMaxBus
    gPciMaxBus = _max

def FindPciDeviceById(vendorId, deviceId, index):
    if drv.gHandle == INVALID_HANDLE_VALUE:
        return 0xFFFFFFFF
    count = 0
    _id = vendorId | (deviceId << 16)
    if vendorId == 0xFFFF:
        return 0xFFFFFFFF
    for bus in range(gPciMaxBus + 1):
        for dev in range(gPciMaxDevice + 1):
            multiFuncFlag = False
            for func in range(gPciMaxFunction + 1):
                if (multiFuncFlag == False) and (func > 0):
                    break
                pciAddress = PciBusDevFunc(bus, dev, func)
                idCode = ReadPciConfigDword(pciAddress, 0)
                if func == 0: # Is Multi Function Device
                    headerType = ReadPciConfigByte(pciAddress, 0x0E)
                    if headerType & 0x80:
                        multiFuncFlag = True
                if idCode == _id:
                    if count == index:
                        return pciAddress
                    count+=1
                    continue
    return 0xFFFFFFFF

def FindPciDeviceByClass(baseClass, subClass, programIf, index):
    if(drv.gHandle == INVALID_HANDLE_VALUE):
        return 0xFFFFFFFF
    count = 0
    _class = (baseClass << 16) | (subClass << 8) | (programIf)
    for bus in range(gPciMaxBus + 1):
        for dev in range(gPciMaxDevice + 1):
            multiFuncFlag = False
            for func in range(gPciMaxFunction + 1):
                if (multiFuncFlag == False) and (func > 0):
                    break
                pciAddress = PciBusDevFunc(bus, dev, func)
                classValue = ReadPciConfigDword(pciAddress, 0x08)
                if func == 0: # Is Multi Function Device
                    headerType = ReadPciConfigByte(pciAddress, 0x0E)
                    if headerType & 0x80:
                        multiFuncFlag = True
                classCode = classValue >> 8 & 0xFFFFFFFF
                if classCode == _class:
                    if count == index:
                        return pciAddress
                    count+=1
                    continue
    return 0xFFFFFFFF
